package com.fuyu.common.utils;


import com.fuyu.common.model.TreeNodeVO;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;

public class TreeUtil {


    /**
     * listToTree
     *
     * @param target      需转换的数据
     * @param getId       主键
     * @param getParentId 父id （父id必须和主键相同类型）
     * @param getChildren 子集
     * @param setChildren 子集
     * @return tree
     */
    public static <T, R> List<T> builder(
            List<T> target,
            Function<T, R> getId,
            Function<T, R> getParentId,
            Function<T, List<T>> getChildren,
            BiConsumer<T, List<T>> setChildren
    ) {
        Map<R, T> oldMap = target.stream().collect(Collectors.toMap(getId, T -> T));
        List<T> result = new ArrayList<>();
        target.forEach(tree -> {
            T parent = oldMap.get(getParentId.apply(tree));
            if (parent == null) {
                result.add(tree);
            } else {
                List<T> ch = getChildren.apply(parent);
                if (ch == null) {
                    ch = new ArrayList<>();
                }
                ch.add(tree);
                setChildren.accept(parent, ch);
            }
        });
        return result;
    }


    public static List<TreeNodeVO> getTreeNodeList(List<TreeNodeVO> list) {
        return TreeUtil.builder(
                list,
                TreeNodeVO::getId,
                TreeNodeVO::getParentId,
                TreeNodeVO::getChildren,
                TreeNodeVO::setChildren
        );
    }

}
